package com.lalameow.researchstation.screen.widget.tree

import com.google.common.cache.CacheBuilder
import com.google.common.cache.CacheLoader
import com.lalameow.researchstation.researching.ResearchingTechnology
import com.lalameow.researchstation.researching.ResearchingTree
import com.lalameow.researchstation.screen.ResearchStationScreen
import com.lalameow.researchstation.screen.widget.tree.ResearchingNodeWidget.Companion.NODE_SIZE
import com.lalameow.researchstation.util.DefaultDirectedGraph
import com.lalameow.researchstation.util.Offset
import com.mojang.blaze3d.systems.RenderSystem
import net.fabricmc.api.EnvType
import net.fabricmc.api.Environment
import net.minecraft.client.font.TextRenderer
import net.minecraft.client.gui.DrawableHelper
import net.minecraft.client.render.BufferRenderer
import net.minecraft.client.render.GameRenderer
import net.minecraft.client.render.VertexFormat.DrawMode
import net.minecraft.client.render.VertexFormats
import net.minecraft.client.util.math.MatrixStack
import net.minecraft.client.util.math.Rect2i
import net.minecraft.util.Identifier
import net.minecraft.util.math.MathHelper
import org.jgrapht.Graph
import org.jgrapht.alg.drawing.CircularLayoutAlgorithm2D
import org.jgrapht.alg.drawing.model.Box2D
import org.jgrapht.alg.drawing.model.LayoutModel2D
import org.jgrapht.alg.drawing.model.MapLayoutModel2D
import org.jgrapht.graph.DefaultEdge
import kotlin.math.floor

typealias TechnologyLayout = LayoutModel2D<ResearchingTechnology>

private fun ResearchingTree.toGraph() = DefaultDirectedGraph<ResearchingTechnology, DefaultEdge>().apply {
    readTree(this@toGraph)
}

private fun Graph<ResearchingTechnology, DefaultEdge>.readTree(tree: ResearchingTree) {
    tree.entries.forEach { (key, value) ->
        addVertex(key)
        value.keys.forEach {
            addVertex(it)
            addEdge(key, it)
        }
        readTree(value)
    }
}

private fun LayoutModel2D<ResearchingTechnology>.layout(graph: Graph<ResearchingTechnology, DefaultEdge>) = apply {
    val algorithm = CircularLayoutAlgorithm2D<ResearchingTechnology, DefaultEdge>(1.0)
    algorithm.layout(graph, this)
}

@Environment(EnvType.CLIENT)
internal class ResearchingTreeWidget(
    tree: ResearchingTree,
    private val contentRect: Rect2i,
    private val textRenderer: TextRenderer,
    private val researched: Set<Identifier>
) : DrawableHelper() {
    companion object {
        const val NODE_VERTICAL_MARGIN = 4
        const val NODE_SIZE_WITH_MARGIN = 30
    }

    private val graph = tree.toGraph()
    private val box =
        Box2D(
            4680 / NODE_SIZE_WITH_MARGIN.toDouble(),
            2260 / NODE_SIZE_WITH_MARGIN.toDouble()
        )
    private val layout = MapLayoutModel2D<ResearchingTechnology>(box).layout(graph)
    private val nodes =
        CacheBuilder.newBuilder().build(object : CacheLoader<ResearchingTechnology, ResearchingNodeWidget>() {
            override fun load(key: ResearchingTechnology) =
                ResearchingNodeWidget(key, textRenderer)
        })

    private val minPos =
        (layout.minOf { it.value.x } - 1) * NODE_SIZE_WITH_MARGIN to
            (layout.minOf { it.value.y } - 1) * NODE_SIZE_WITH_MARGIN
    private val maxPos =
        (layout.maxOf { it.value.x } + 1) * NODE_SIZE_WITH_MARGIN to
            (layout.maxOf { it.value.y } + 1) * NODE_SIZE_WITH_MARGIN
    var currentPos = 0 to 0

    fun MatrixStack.draw(
        screenX: Int,
        screenY: Int,
        mouse: Offset
    ) = drawLayout(layout, screenX, screenY, mouse)

    private fun MatrixStack.drawLayout(
        layout: TechnologyLayout,
        screenX: Int,
        screenY: Int,
        mouse: Offset
    ) {
        push()
        val tessellator = RenderSystem.renderThreadTesselator()
        val bufferBuilder = tessellator.buffer
        val matrix = peek().positionMatrix

        RenderSystem.disableTexture()
        RenderSystem.setShader {
            GameRenderer.getPositionColorShader()!!.also { shader ->
                shader.lineWidth?.set(4.0f)
            }
        }
        RenderSystem.lineWidth(4.0f)
        graph.edgeSet().forEach {
            val sourcePoint = layout[graph.getEdgeSource(it)]
            val targetPoint = layout[graph.getEdgeTarget(it)]
            val sourceX =
                currentPos.first + ResearchStationScreen.SIDE + sourcePoint.x * NODE_SIZE_WITH_MARGIN + NODE_SIZE / 2 - minPos.first
            val sourceY =
                currentPos.second + ResearchStationScreen.TOP + sourcePoint.y * NODE_SIZE_WITH_MARGIN + NODE_SIZE / 2 - minPos.second
            val targetX =
                currentPos.first + ResearchStationScreen.SIDE + targetPoint.x * NODE_SIZE_WITH_MARGIN + NODE_SIZE / 2 - minPos.first
            val targetY =
                currentPos.second + ResearchStationScreen.TOP + targetPoint.y * NODE_SIZE_WITH_MARGIN + NODE_SIZE / 2 - minPos.second

            bufferBuilder.begin(DrawMode.DEBUG_LINES, VertexFormats.POSITION_COLOR)
            bufferBuilder.vertex(matrix, floor(sourceX).toFloat(), floor(sourceY).toFloat(), 0.0f)
                .color(0, 0, 0, 255)
                .next()
            bufferBuilder.vertex(matrix, floor(targetX).toFloat(), floor(targetY).toFloat(), 0.0f)
                .color(0, 0, 0, 255)
                .next()
            bufferBuilder.end()
            BufferRenderer.draw(bufferBuilder)
        }
        RenderSystem.enableTexture()
        RenderSystem.lineWidth(1.0f)

        layout.forEach { (technology, point) ->
            val node = nodes[technology]
            with(node) {
                val x = MathHelper.floor(ResearchStationScreen.SIDE + point.x * NODE_SIZE_WITH_MARGIN - minPos.first)
                val y = MathHelper.floor(ResearchStationScreen.TOP + point.y * NODE_SIZE_WITH_MARGIN - minPos.second)
                draw(
                    currentPos.first + x to currentPos.second + y,
                    technology.id in researched,
                    screenX,
                    screenY,
                    mouse
                )
            }
        }
        pop()
    }

    fun move(deltaX: Double, deltaY: Double) {
        val x = MathHelper.floor(deltaX + currentPos.first)
        val y = MathHelper.floor(deltaY + currentPos.second)
        if (
            x > 0 && x + contentRect.width < maxPos.first - minPos.first &&
            y > 0 && y + contentRect.height < maxPos.second - minPos.second
        ) currentPos = x to y
    }
}
